import json
import re
from typing import List, Iterator, Dict, Any
from .medical_record import MedicalRecord


class MedicalRecordIterator:
    """
    电子病历迭代器类
    实现迭代器模式，支持逐批读取病历数据，避免一次性加载所有数据导致内存溢出
    """
    
    def __init__(self, records_data: List[Dict[str, Any]], batch_size: int = 10):
        """
        初始化迭代器
        
        Args:
            records_data (List[Dict[str, Any]]): 病历数据列表，每个元素是一个包含病历信息的字典
            batch_size (int): 每次迭代返回的病历数量，默认为10
        """
        self.records_data = records_data
        self.batch_size = batch_size
        self.current_index = 0
    
    def __iter__(self) -> Iterator[List[MedicalRecord]]:
        """
        返回迭代器对象本身
        
        Returns:
            Iterator[List[MedicalRecord]]: 迭代器对象
        """
        return self
    
    def __next__(self) -> List[MedicalRecord]:
        """
        返回下一批病历数据
        
        Returns:
            List[MedicalRecord]: 下一批病历数据列表
            
        Raises:
            StopIteration: 当没有更多病历数据时抛出
        """
        # 检查是否还有未处理的病历数据
        if self.current_index >= len(self.records_data):
            raise StopIteration
        
        # 计算当前批次的结束索引
        end_index = min(self.current_index + self.batch_size, len(self.records_data))
        
        # 提取当前批次的病历数据
        batch_data = self.records_data[self.current_index:end_index]
        
        # 将字典数据转换为MedicalRecord对象
        batch_records = []
        for record_data in batch_data:
            # 清理数据中的多余空格和换行符
            patient_id = self._clean_text(record_data.get('patient_id', ''))
            diagnosis = self._clean_text(record_data.get('diagnosis', ''))
            medications = record_data.get('medications', [])
            report_date = self._clean_text(record_data.get('report_date', ''))
            
            # 如果medications是字符串，尝试解析为列表
            if isinstance(medications, str):
                try:
                    medications = json.loads(medications)
                except json.JSONDecodeError:
                    # 如果解析失败，将其作为单个元素的列表处理
                    medications = [medications] if medications else []
            
            record = MedicalRecord(
                patient_id=patient_id,
                diagnosis=diagnosis,
                medications=medications,
                report_date=report_date,
                raw_data=json.dumps(record_data)
            )
            batch_records.append(record)
        
        # 更新当前索引
        self.current_index = end_index
        
        return batch_records
    
    def _clean_text(self, text: str) -> str:
        """
        清理文本中的多余空格和换行符
        
        Args:
            text (str): 原始文本
            
        Returns:
            str: 清理后的文本
        """
        if not isinstance(text, str):
            return str(text)
        # 移除首尾空格和换行符，并将多个连续空格替换为单个空格
        return re.sub(r'\s+', ' ', text.strip())


# 示例用法（可选）
if __name__ == "__main__":
    # 模拟病历数据
    sample_data = [
        {
            "patient_id": "A12345678",
            "diagnosis": "高血压",
            "medications": ["降压药A", "降压药B"],
            "report_date": "2023-05-15"
        },
        {
            "patient_id": "B87654321",
            "diagnosis": "糖尿病",
            "medications": ["胰岛素"],
            "report_date": "2023-05-16"
        }
    ]
    
    # 创建迭代器并测试
    iterator = MedicalRecordIterator(sample_data, batch_size=1)
    for i, batch in enumerate(iterator):
        print(f"批次 {i+1}: {batch}")